import os
from sqlalchemy import create_engine
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import sessionmaker, Session
from typing import AsyncGenerator
import config

# 数据库配置
DATABASE_TYPE = getattr(config, 'DATABASE_TYPE', 'sqlite')  # sqlite 或 mysql
DATABASE_URL = getattr(config, 'DATABASE_URL', None)

if DATABASE_URL is None:
    if DATABASE_TYPE == 'sqlite':
        DATABASE_URL = "sqlite:///./capcut_api.db"
        ASYNC_DATABASE_URL = "sqlite+aiosqlite:///./capcut_api.db"
    else:  # mysql
        DATABASE_URL = "mysql+pymysql://root:password@localhost:3306/capcut_api"
        ASYNC_DATABASE_URL = "mysql+aiomysql://root:password@localhost:3306/capcut_api"
else:
    # 如果提供了DATABASE_URL，自动生成异步版本
    if DATABASE_URL.startswith('sqlite:///'):
        ASYNC_DATABASE_URL = DATABASE_URL.replace('sqlite:///', 'sqlite+aiosqlite:///')
    elif DATABASE_URL.startswith('mysql+pymysql://'):
        ASYNC_DATABASE_URL = DATABASE_URL.replace('mysql+pymysql://', 'mysql+aiomysql://')
    else:
        ASYNC_DATABASE_URL = DATABASE_URL

# 创建同步引擎（用于Alembic迁移）
engine = create_engine(
    DATABASE_URL,
    echo=getattr(config, 'DATABASE_ECHO', False),
    pool_pre_ping=True,
    pool_recycle=300 if DATABASE_TYPE == 'mysql' else -1
)

# 创建异步引擎（用于FastAPI异步操作）
async_engine = create_async_engine(
    ASYNC_DATABASE_URL,
    echo=getattr(config, 'DATABASE_ECHO', False),
    pool_pre_ping=True,
    pool_recycle=300 if DATABASE_TYPE == 'mysql' else -1
)

# 创建会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
AsyncSessionLocal = async_sessionmaker(
    async_engine,
    class_=AsyncSession,
    expire_on_commit=False
)

# 依赖注入：获取数据库会话
def get_db() -> Session:
    """获取同步数据库会话"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

async def get_async_db() -> AsyncGenerator[AsyncSession, None]:
    """获取异步数据库会话"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close() 